#ifndef _T_BUFFER_2D_CU_H_
#define _T_BUFFER_2D_CU_H_

//////////////////////////////////////////////////////////////////////////
/// T_cudaBuffer2D is an interface for cuda mappable objects 
/// like host memory or d3d ogl devices
//////////////////////////////////////////////////////////////////////////
template <class T>
class T_Buffer2D_CU : public I_Buffer2D, public T_Buffer_CU<T>
{
public:
	T_Buffer2D_CU(void) {}

	virtual ~T_Buffer2D_CU(void) { Release(); };

	//////////////////////////////////////////////////////////////////////////
	/// I_Buffer2D knows how to create and release its pointers
	//////////////////////////////////////////////////////////////////////////
	virtual void Create(int32 nWidth, int32 nHeight, int32 nPitch, I_api_Device *pDevice = NULL)
	{
		if (!nWidth || !nHeight) return;

		m_width = nWidth;
		m_height = nHeight;
		m_pitch = nPitch;

		m_size = (nWidth + nPitch) * nHeight;

		m_hptr = (T *) malloc(SizeInBytes());

		cudaMalloc((void **) &m_dptr, SizeInBytes());

		m_offset = 0;

		m_owner = true;
	}

	virtual void Release(void)
	{
		if (m_owner) 
		{
			m_pBuffer->Release();
		}

		m_hptr = NULL;
		m_dptr = NULL;
		m_flag = 0;
		m_size = m_width = m_height = m_pitch = 0;

		m_owner = false;
	}

	//////////////////////////////////////////////////////////////////////////
	virtual void *Map(int32 nFlag)
	{
		// if the buffer has already been mapped return NULL instantly
		if (m_flag) return NULL;

		// otherwise act as in e_Map description
		switch(nFlag)
		{
			///  When you need to get data from cuda pointer
			///    Call Map(RO)
			///    Map copies data from cuda pointer to host and returns host pointer
			///    Call Unmap()
			///    Read host pointer
			///    Unmap() doesn't change anything
		case e_Map_RO: 
			cudaMemcpy(m_hptr, m_dptr, SizeInBytes(), cudaMemcpyDeviceToHost);
			break;
			///  When you need to write initial values to cuda pointer
			///    Call Map(WO)
			///    Map instantly returns host pointer
			///    Write into host pointer
			///    Call Unmap ()
			///    Unmap checks that flag was WO and writes back to cuda array
		case e_Map_WO:
			break;
			///  When you need to Read cuda pointer Modify and Write
			///    Call Map(RW)
			///    Map copies data from cuda pointer to host and returns host pointer
			///    Read/Write host pointer 
			///    Call Unmap()
			///    Unmap() copies data from host pointer to cuda pointer
		case e_Map_RW:
			cudaMemcpy(m_hptr, m_dptr, SizeInBytes(), cudaMemcpyDeviceToHost);
			break;
			///  In some cases what you need is an OGL or D3D buffer and you want to 
			///  map pointers to CUDA. the way you do it is 
			///    Call Map(BO)
			///    Map return device pointer
			///    Call Unmap
			///    Unmap does unmapping if necessary
		case e_Map_Lock:
			m_flag = nFlag;
			return m_dptr; 
			break;
		default: return NULL;
		}

		// save flag for subsequent unmap call
		m_flag = nFlag;

		// map always returns host pointer
		return m_hptr;
	}

	virtual void Unmap()
	{
		// otherwise act as in e_Map description
		switch(m_flag)
		{
			///  When you need to get data from cuda pointer
			///    Call Map(RO)
			///    Map copies data from cuda pointer to host and returns host pointer
			///    Call Unmap()
			///    Read host pointer
			///    Unmap() doesn't change anything
		case e_Map_RO: 

			break;
			///  When you need to write initial values to cuda pointer
			///    Call Map(WO)
			///    Map instantly returns host pointer
			///    Write into host pointer
			///    Call Unmap ()
			///    Unmap checks that flag was WO and writes back to cuda array
		case e_Map_WO:
			cudaMemcpy(m_dptr, m_hptr, SizeInBytes(), cudaMemcpyHostToDevice);
			break;
			///  When you need to Read cuda pointer Modify and Write
			///    Call Map(RW)
			///    Map copies data from cuda pointer to host and returns host pointer
			///    Read/Write host pointer 
			///    Call Unmap()
			///    Unmap() copies data from host pointer to cuda pointer
		case e_Map_RW:
			cudaMemcpy(m_dptr, m_hptr, SizeInBytes(), cudaMemcpyHostToDevice);
			break;
		case e_Map_Lock:
			// Do nothing
			break;

		default: 
			// if anything else was called return now to avoid reseting the flag
			return;
		}

		// reset flag
		m_flag = 0;
	}

	//////////////////////////////////////////////////////////////////////////
	/// I_Buffer2D knows how to copy data back and forth
	/// I_Buffer2D allows to some way to exchange data between pointers
	/// without any notion of what those pointers really are
	/// Copy direction DtoH copies from start index to stop device to host
	/// Copy direction HtoD copies from start index to stop host   to device
	//////////////////////////////////////////////////////////////////////////
	virtual void Memcpy(int32 dir, int32 nStart = 0 , int32 nCount = 0)
	{
		if (nCount == 0) nCount = m_size-nStart;

		switch(dir)
		{
		case e_Copy_DtoH:
			cudaMemcpy(m_hptr+nStart, m_dptr+nStart, nCount*sizeof(T), cudaMemcpyDeviceToHost);
			break;
		case e_Copy_HtoD:
			cudaMemcpy(m_dptr+nStart, m_hptr+nStart, nCount*sizeof(T), cudaMemcpyHostToDevice);
			break;

		default: 
			break;
		}
	}

	virtual void Memset(void *ptr, int32 nStart = 0 , int32 nCount = 0)
	{
		if (nCount == 0) nCount = m_size-nStart;

		T t;

		if (ptr == NULL)
			memset(&t, 0, sizeof(T));
		else
			memcpy(&t, ptr, sizeof(T));

		for (int i = nStart; i < nStart + nCount; i++)
			m_hptr[i] = t;

		cudaMemcpy(m_dptr+nStart, m_hptr+nStart, nCount*sizeof(T), cudaMemcpyHostToDevice);
	}

	virtual void				Memcpy(int32 nDir, 
														 int32 nStartR = 0, int32 nStartC = 0,
														 int32 nStopR = 0, int32 nStopC = 0)
	{
		if (nStopC == 0) nStopC = m_width-nStartC;
		if (nStopR == 0) nStopR = m_height-nStartR;

		switch(nDir)
		{
		case e_Copy_DtoH:
			for (int ir = nStartR; ir < nStopR; ir++)
				cudaMemcpy(m_hptr+ir*(m_width+m_pitch)+nStartC, 
									 m_dptr+ir*(m_width+m_pitch)+nStartC, 
									 (nStopC - nStartC)*sizeof(T), 
									 cudaMemcpyDeviceToHost);
			break;
		case e_Copy_HtoD:
			for (int ir = nStartR; ir < nStopR; ir++)
				cudaMemcpy(m_hptr+ir*(m_width+m_pitch)+nStartC, 
									 m_dptr+ir*(m_width+m_pitch)+nStartC, 
									 (nStopC - nStartC)*sizeof(T), 
									 cudaMemcpyHostToDevice);
			break;

		default: 
			break;
		}
	}

	virtual void				Memset(void *pElem, 
														 int32 nStartR = 0, int32 nStartC = 0,
														 int32 nStopR = 0, int32 nStopC = 0)
	{
		if (nStopC == 0) nStopC = m_width-nStartC;
		if (nStopR == 0) nStopR = m_height-nStartR;

		T t;

		if (pElem == NULL)
			memset(&t, 0, sizeof(T));
		else
			memcpy(&t, pElem, sizeof(T));

		for (int ir = nStartR; ir < nStopR; ir++)
		{
			// set values within a row
			for (int ic = nStartC; ic < nStopC; ic++)
				m_hptr[nStartC+ic+ir*(m_width+m_pitch)] = t;

			// copy the whole row
			cudaMemcpy(m_hptr+ir*(m_width+m_pitch)+nStartC, 
								 m_dptr+ir*(m_width+m_pitch)+nStartC, 
								 (nStopC - nStartC)*sizeof(T), 
								 cudaMemcpyHostToDevice);
		}
	}

	virtual int					SizeInBytes(void) { return m_size * sizeof(T); }
};

#endif